package org.raymond.iworks.study.mq.rabbitmq.classic.pubcon;

import com.rabbitmq.client.Channel;
import com.rabbitmq.client.ConfirmCallback;
import com.rabbitmq.client.Connection;
import com.rabbitmq.client.ConnectionFactory;
import lombok.extern.slf4j.Slf4j;

import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.util.UUID;
import java.util.concurrent.ConcurrentNavigableMap;
import java.util.concurrent.ConcurrentSkipListMap;
import java.util.function.BooleanSupplier;

/**
 * @author raymond
 * @version V1.0
 * @Description: TODO
 */
@Slf4j
public class PublisherConfirms {
    private static final int MESSAGE_COUNT = 50_000;
    public static void main(String[] args) throws Exception {
        publishingMessagesIndividually();
        publishingMessagesInBatches();
        handlePublishConfirmsAsynchronously();
    }

    // 发布消息并同步等待其确认
    public static void publishingMessagesIndividually() throws Exception {
        try(Connection conn = createConnnection();
            Channel channel = conn.createChannel()){
            String queueName = UUID.randomUUID().toString();
            boolean durable = false;
            channel.queueDeclare(queueName, durable, false, true, null);

            // 必须在期望使用发布者确认的每个通道上调用此方法.
            // 确认应该只启用一次,而不是对每条发布的消息都启用
            channel.confirmSelect();
            long start = System.nanoTime();
            for(int i=0; i<MESSAGE_COUNT; i++){
                String body = String.valueOf(i);
                channel.basicPublish("", queueName, null, body.getBytes(StandardCharsets.UTF_8));
                // 以同步方式等待,确认消息后,该方法将立即返回.如果消息未在超时内得到确认,或者消息被破坏了(意味着代理由于某种原因无法处理它),
                // 该方法将引发异常.异常的处理通常包括记录错误消息和/或重试发送消息
                // 不同的客户端有不同的方式来同步处理发布者确认
                // 缺点:大大减慢了发布速度.因为消息的确认会阻止所有后续消息的发布.这种方法提供的吞吐量不会超过每秒几百条
                channel.waitForConfirmsOrDie(5_000);
            }
            long end = System.nanoTime();
            log.info("Published {} message individually in {}",  MESSAGE_COUNT, Duration.ofNanos(end - start).toMillis());
        }
    }

    public static void publishingMessagesInBatches() throws Exception {
        try(Connection conn = createConnnection();
            Channel channel = conn.createChannel()){
            String queueName = UUID.randomUUID().toString();
            boolean durable = false;
            channel.queueDeclare(queueName, durable, false, true, null);

            // 必须在期望使用发布者确认的每个通道上调用此方法.
            // 确认应该只启用一次,而不是对每条发布的消息都启用
            channel.confirmSelect();

            // 等待整个批次得到确认
            int batchSize = 100;
            int outstandingMessageCount = 0;

            long start = System.nanoTime();
            for(int i=0; i<MESSAGE_COUNT; i++){
                String body = String.valueOf(i);
                channel.basicPublish("", queueName, null, body.getBytes(StandardCharsets.UTF_8));
                outstandingMessageCount++;

                if(outstandingMessageCount == batchSize){
                    // 依然是同步的,等待这一批消息确认时依然会阻塞后续消息的发布
                    // 发生问题时不知道哪里出了问题,需要保留整个批处理有意义的内容或者重新发布
                    channel.waitForConfirmsOrDie(5_000);
                    outstandingMessageCount = 0;
                }
            }

            if(outstandingMessageCount>0){
                channel.waitForConfirmsOrDie(5_000);
            }
            long end = System.nanoTime();
            log.info("Published {} message in batch in {}",  MESSAGE_COUNT, Duration.ofNanos(end - start).toMillis());
        }
    }

    public static void handlePublishConfirmsAsynchronously() throws Exception {
        try(Connection conn = createConnnection();
            Channel channel = conn.createChannel()){
            String queueName = UUID.randomUUID().toString();
            boolean durable = false;
            channel.queueDeclare(queueName, durable, false, true, null);

            // 必须在期望使用发布者确认的每个通道上调用此方法.
            // 确认应该只启用一次,而不是对每条发布的消息都启用
            channel.confirmSelect();

            // 异步发布确认的步骤-1. 提供一种将publishSeqNo与消息关联起来的方法
            // 使用map将sequenceNumber与request body关联起来
            ConcurrentNavigableMap<Long, String> outstandingConfirms = new ConcurrentSkipListMap();
            // when message is confirmed
            // sequenceNumber: 用来标识已确认或否认的消息
            // multiple: false:只确认一条消息,true:则所有<=sequenceNumber的消息都被确认/否认
            ConfirmCallback cleanWhenAckCallback = (sequenceNumber, multiple)->{
                if(multiple){
                    ConcurrentNavigableMap<Long, String> confirmed = outstandingConfirms.headMap(
                            sequenceNumber, true
                    );
                    confirmed.clear();
                }else{
                    outstandingConfirms.remove(sequenceNumber);
                }
                // 不要尝试在回调中重新发布否认的消息,因为确认回调是在通道不应该执行操作的IO线程中分配的
                // 最好是在类似ConcurrentLinkedQueue这样的内存队列中对消息进行排队,由发布线程轮询
            };

            // when message is nack-ed
            ConfirmCallback nackCallback = (sequenceNumber, multiple)->{
                String body = outstandingConfirms.get(sequenceNumber);
                log.error("Message with body {} has been nack-ed. Sequence number: {}, multiple: {}",
                        body, sequenceNumber, multiple);
                cleanWhenAckCallback.handle(sequenceNumber, multiple);
            };

            // 异步发布确认的步骤-2. 在客户端上注册一个回调就可以收到确认/否认的通知
            channel.addConfirmListener(cleanWhenAckCallback, nackCallback);

            long start = System.nanoTime();
            for(int i=0; i<MESSAGE_COUNT; i++){
                String body = String.valueOf(i);
                // 异步发布确认的步骤-3. 在发布消息之前跟踪publishSeqNo
                outstandingConfirms.put(channel.getNextPublishSeqNo(), body);
                channel.basicPublish("", queueName, null, body.getBytes(StandardCharsets.UTF_8));
            }

            if(!waitUnit(Duration.ofSeconds(60), outstandingConfirms::isEmpty)){
                throw new IllegalStateException("All messages could not be confirmed in 60 seconds");
            }
            long end = System.nanoTime();
            log.info("Published {} message in batch in {}",  MESSAGE_COUNT, Duration.ofNanos(end - start).toMillis());
        }
    }

    private static boolean waitUnit(Duration timeout, BooleanSupplier condition) throws InterruptedException {
        int waited = 0;
        while(!condition.getAsBoolean() && waited < timeout.toMillis()){
            Thread.sleep(100L);
            waited += 100;
        }
        return condition.getAsBoolean();
    }

    private static Connection createConnnection() throws Exception {
        ConnectionFactory cf = new ConnectionFactory();
        cf.setHost("192.168.0.125");
        cf.setPort(5672);
        return cf.newConnection();
    }
}
